import json
from typing import List

from user import *


def get_message(chat_history: List[list], docs: List[str], query: str):
    result = []
    history = '\n'.join([item['role'] + ': ' + item['content'] for item in chat_history])
    doc = '\n'.join([d.page_content for d in docs])
    result.append(
        {
            'role': 'system',
            'content': '你是一个用户友好的助手， 会根据参考信息和沟通历史回答用户的问题。\n'
                       '沟通历史： \n'
                       f'{history}\n'
                       '参考信息：\n'
                       f'{doc}'
        }
    )
    result.append(
        {
            "role": "user",
            "content": query
        }
    )
    return result


def get_response_stream_generate_from_ChatGPT_API(
        message_context,
        apikey,
        message_history,
        model="gpt-3.5-turbo",
        temperature=0.9,
        presence_penalty=0,
        max_tokens=2000):
    """
    从ChatGPT API获取回复
    :param apikey:
    :param message_context: 上下文
    :param message_history: 消息历史
    :param model: 模型
    :param temperature: 温度
    :param presence_penalty: 惩罚
    :param max_tokens: 最大token数量
    :return: 回复生成器
    """
    if apikey is None:
        apikey = API_KEY
    header = {
        "Content-Type": "application/json",
        "Authorization": "Bearer " + apikey
    }

    data = {
        "model": model,
        "temperature": temperature,
        "presence_penalty": presence_penalty,
        "max_tokens": max_tokens,
        "messages": message_context,
        "stream": True
    }
    logger.info("开始流式请求")
    url = f"{API_BASE}/chat/completions"
    # 请求接收流式数据 动态print
    try:
        print(f"request body:{json.dumps(data, ensure_ascii=False, indent=2)}")
        response = requests.request("POST", url, headers=header, json=data, stream=True)
        print(f"response body:{json.dumps(response, ensure_ascii=False, indent=2)}")

        def generate():
            stream_content = str()
            one_message = {"role": "assistant", "content": stream_content}
            message_history.append(one_message)
            i = 0
            for line in response.iter_lines():
                # print(str(line))
                line_str = str(line, encoding='utf-8')
                if line_str.startswith("data:"):
                    if line_str.startswith("data: [DONE]"):
                        asyncio_run(save_all_user_dict())
                        logger.info("用户得到的回复内容：{}...".format(one_message["content"][:200]))
                        break
                    line_json = json.loads(line_str[5:])
                    if 'choices' in line_json:
                        if len(line_json['choices']) > 0:
                            choice = line_json['choices'][0]
                            if 'delta' in choice:
                                delta = choice['delta']
                                if 'role' in delta:
                                    role = delta['role']
                                elif 'content' in delta:
                                    delta_content = delta['content']
                                    i += 1
                                    if i < 40:
                                        print(delta_content, end="")
                                    elif i == 40:
                                        print("......")
                                    one_message['content'] = one_message['content'] + delta_content
                                    yield delta_content

                elif len(line_str.strip()) > 0:
                    print(line_str)
                    yield line_str

    except Exception as e:
        ee = e

        def generate():
            yield "request error:\n" + str(ee)

    return generate


def get_response_from_ChatGPT_API(
        message_context,
        apikey,
        docs,
        query,
        model="gpt-3.5-turbo",
        temperature=0.9,
        presence_penalty=0,
        max_tokens=2000):
    """
    从ChatGPT API获取回复
    :param message_context: 上下文
    :param apikey: API KEY
    :param model: 模型
    :param temperature: 温度
    :param presence_penalty: 惩罚
    :param max_tokens: 最大token数量
    :return: 回复
    """
    if apikey is None:
        apikey = API_KEY

    header = {
        "Content-Type": "application/json",
        "Authorization": "Bearer " + apikey
    }

    message = get_message(message_context, docs, query)
    print(json.dumps(message, indent=2, ensure_ascii=False))
    data = {
        "model": model,
        "messages": message,
        "temperature": temperature,
        "presence_penalty": presence_penalty,
        "max_tokens": max_tokens
    }
    url = f"{API_BASE}/chat/completions"

    try:
        print(f"request body:{json.dumps(data, ensure_ascii=False, indent=2)}")
        response = requests.post(url, headers=header, data=json.dumps(data))
        print(response)
        # print(f"response body:{json.dumps(response, ensure_ascii=False, indent=2)}")
        response = response.json()
        # 判断是否含 choices[0].message.content
        if "choices" in response \
                and len(response["choices"]) > 0 \
                and "message" in response["choices"][0] \
                and "content" in response["choices"][0]["message"]:
            data = response["choices"][0]["message"]["content"]
        else:
            data = str(response)

    except Exception as e:
        logger.error(e)
        return str(e)

    return data


def get_balance(apikey):
    head = ""
    if apikey is not None:
        head = "###  用户专属api key余额  \n"
    else:
        head = "### 通用api key  \n"
        apikey = API_KEY

    subscription_url = f"{API_BASE}/dashboard/billing/subscription"
    headers = {
        "Authorization": "Bearer " + apikey,
        "Content-Type": "application/json"
    }
    subscription_response = requests.get(subscription_url, headers=headers)
    if subscription_response.status_code == 200:
        data = subscription_response.json()
        total = data.get("hard_limit_usd")
    else:
        return head + subscription_response.text

    # start_date设置为今天日期前99天
    start_date = (datetime.datetime.now() - datetime.timedelta(days=99)).strftime("%Y-%m-%d")
    # end_date设置为今天日期+1
    end_date = (datetime.datetime.now() + datetime.timedelta(days=1)).strftime("%Y-%m-%d")
    billing_url = f"{API_BASE}/dashboard/billing/usage?start_date={start_date}&end_date={end_date}"
    billing_response = requests.get(billing_url, headers=headers)
    if billing_response.status_code == 200:
        data = billing_response.json()
        total_usage = data.get("total_usage") / 100
        daily_costs = data.get("daily_costs")
        days = min(5, len(daily_costs))
        recent = f"##### 最近{days}天使用情况  \n"
        for i in range(days):
            cur = daily_costs[-i - 1]
            date = datetime.datetime.fromtimestamp(cur.get("timestamp")).strftime("%Y-%m-%d")
            line_items = cur.get("line_items")
            cost = 0
            for item in line_items:
                cost += item.get("cost")
            recent += f"\t{date}\t{cost / 100} \n"
    else:
        return head + billing_response.text

    return head + f"\n#### 总额:\t{total:.4f}  \n" \
                  f"#### 已用:\t{total_usage:.4f}  \n" \
                  f"#### 剩余:\t{total - total_usage:.4f}  \n" \
                  f"\n" + recent
